# -*- coding: utf-8 -*-
"""
Created on Sun Aug 29 09:28:00 2021

@author: Administrator
"""

import tensorflow as tf
from tensorflow.keras import datasets, layers, optimizers, Sequential, metrics
from tensorflow import keras

def process( x,y ):
    '''
    Parameters
    ----------
    x : 是一个单独的图片，不是batch
        DESCRIPTION.
    y : TYPE
        DESCRIPTION.
    '''
    x = tf.cast( x, dtype=tf.float32 ) / 255.0;
    y = tf.cast( y, dtype = tf.int32 );
    return x,y;



# 加载数据，以及简单的预处理
batchsize = 128;
(x,y), (x_val, y_val) = datasets.cifar10.load_data();
y = tf.squeeze(y);
y_val = tf.squeeze( y_val );

y = tf.one_hot(y, depth = 10);
y_val = tf.one_hot( y_val, depth = 10 );

print( 'datasets:', x.shape, y.shape, x.min(), x.max() );

# 构建训练数据集
train_db = tf.data.Dataset.from_tensor_slices( (x,y) );
train_db = train_db.map( process ).shuffle(10000).batch( batchsize );

test_db = tf.data.Dataset.from_tensor_slices( (x_val, y_val) );
test_db = test_db.map( process ).batch( batchsize );

print( "data base:", type( train_db )  )


sample = next( iter(train_db) );
print( "sample:",  type(sample), sample[0].shape, sample[1].shape )


# 自定义一个层，是一个类
class MyDense( layers.Layer ):
    def __init__(self, input_dim, output_dim):
        super(MyDense, self).__init__();
        
        self.kernel = self.add_variable( 'w',[input_dim, output_dim] )
        #self.bias = self.add_variable( 'b', [output_dim] );
        
    def call(self, inputs, training = None):
        x = inputs @ self.kernel;
        return x;

# 自定义一个网络
class MyNetwork( keras.Model ):
    def __init__(self):
        super( MyNetwork, self ).__init__();
        
        self.fc1 = MyDense( 32*32*3, 256);
        self.fc2 = MyDense( 256, 128);
        self.fc3 = MyDense( 128, 64 );
        self.fc4 = MyDense( 64, 32 );
        self.fc5 = MyDense( 32, 10 );
        
    def call( self, inputs, training = None ):
        # 实现前向传播的逻辑
        x = tf.reshape( inputs,[-1,32*32*3] );
        x = self.fc1(x);
        x = tf.nn.relu( x );
        
        x = self.fc2(x);
        x = tf.nn.relu( x );
        
        x = self.fc3(x);
        x = tf.nn.relu( x );
        
        x = self.fc4(x);
        x = tf.nn.relu( x );
        
        x = self.fc5(x);
        #x = tf.nn.relu( x );
        
        return x;



network = MyNetwork();


# "装配起来"
network.compile( optimizer = optimizers.Adam(lr = 1e-4), 
                loss = tf.losses.CategoricalCrossentropy(from_logits=True),
                metrics = ['accuracy']);

network.fit( train_db,epochs = 5, 
            validation_data = test_db, validation_freq = 1 )

# 保存模型
network.evaluate( test_db );
network.save_weights( 'ckpt/weights.ckpt' );
del network;

print( 'saved' );

network = MyNetwork();
network.compile( optimizer = optimizers.Adam(lr = 1e-4), 
                loss = tf.losses.CategoricalCrossentropy(from_logits=True),
                metrics = ['accuracy']);
network.load_weights( 'ckpt/weights.ckpt' );
network.evaluate( test_db );

